#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import ntpath
import time
import argparse as ap
import requests
from tqdm import tqdm
from urllib.parse import quote
from loguru import logger
import concurrent.futures
from datetime import datetime
import re
from typing import Dict, Any, Optional, List


class ImageDownloader:
    def __init__(self,
                 # 通用
                 startdate: str,
                 enddate: str,
                 focus: bool = False,
                 ell_upper_limit: Optional[float] = None,
                 fwhm_upper_limit: Optional[float] = None,
                 workers: int = 4,
                 api_version: str = "V20250303_KM",
                 local: bool = False,
                 au_key: str = "",
                 bands: Optional[List[str]] = None,
                 # object 模式
                 targetname: Optional[str] = None,
                 # radec 模式
                 ra_min: Optional[float] = None,
                 ra_max: Optional[float] = None,
                 dec_min: Optional[float] = None,
                 dec_max: Optional[float] = None,
                 ):
        """
        一个实例支持两种模式：
          1) object 模式：传入 targetname 或 target list 文件
          2) RA/DEC 框选模式：传入 ra_min, ra_max, dec_min, dec_max
        自动根据入参判断模式。
        """
        self.targetname = targetname
        self.local = local
        self.startdate = startdate
        self.enddate = enddate
        self.focus = focus
        self.ell_upper_limit = ell_upper_limit
        self.fwhm_upper_limit = fwhm_upper_limit
        self.workers = max(1, int(workers))
        self.api_version = api_version
        self.au_key = au_key

        self.bands = bands if bands else ["u", "v", "g", "r", "i", "z"]

        # RA/DEC box
        self.ra_min = ra_min
        self.ra_max = ra_max
        self.dec_min = dec_min
        self.dec_max = dec_max

        # 判断模式
        self.mode = self._decide_mode()

        # 接口地址
        self.api_host = "http://192.168.16.70:8082"
        if self.mode == "object":
            self.base_url = f"{self.api_host}/{self.api_version}/search_by_using_obj_filter"
        else:
            # radec
            self.base_url = f"{self.api_host}/{self.api_version}/search_by_radec_box"

        # 本地路径前缀（仅 local=True 时使用）
        self.localurl = "/all/data_products/"

        # requests Session（可统一设置重试或headers）
        self.session = requests.Session()

    # ------------------------ 模式与参数 ------------------------

    def _decide_mode(self) -> str:
        """根据入参判断使用 object 还是 radec 模式。"""
        has_radec = all(v is not None for v in [self.ra_min, self.ra_max, self.dec_min, self.dec_max])
        if has_radec:
            return "radec"
        # 兼容你的原脚本：targetname 存在就是 object 模式（可为单个名称或文件）
        if self.targetname:
            return "object"
        raise ValueError("请至少提供 object（targetname）或 RA/DEC 边界（ra_min/ra_max/dec_min/dec_max）之一。")

    # ------------------------ URL 映射（local 模式） ------------------------

    def transform_urls(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """转换URL中的端口到对应的存储节点路径（仅在 local=True 时使用）。"""
        transformed = data.copy()

        if self.api_version == "V20250303_KM":
            # 生成所有 1..29 节点的映射规则
            port_mappings = {}
            for node_id in range(1, 30):
                # 端口号 10001 到 10029
                port = f"100{node_id:02d}"
                node_name = f"kmnode{node_id:02d}"
                pattern = re.compile(rf'http://[^:/]+:{port}')
                replacement = f"/all/data_products/{node_name}/1m6"
                port_mappings[port] = (replacement, pattern)

            for key, value in transformed.items():
                if isinstance(value, str):
                    current = value
                    for port, (replacement, pattern) in port_mappings.items():
                        current = pattern.sub(replacement, current)
                    transformed[key] = current

            return transformed

        elif self.api_version.startswith("V20250303"):
            for key, value in transformed.items():
                if isinstance(value, str):
                    current = value
                    # 1. 固定 IP 的 16000 端口（总汇）
                    current = re.sub(
                        r"http://192\.168\.16\.206:16000/V20250303",
                        f"/home/mpilot/1m6/reception/V20250303",
                        current
                    )
                    # 2. 每个节点的 reception 数据（10001~10029）
                    for node_id in range(1, 30):
                        port = f"100{node_id:02d}"
                        pattern = re.compile(rf"http://[^:/]+:{port}/reception/V20250303")
                        replacement = f"/all/data_products/kmnode{node_id:02d}/1m6/reception/V20250303"
                        current = pattern.sub(replacement, current)
                    transformed[key] = current
            return transformed

        else:
            raise ValueError(f"Unsupported API version for local mapping: {self.api_version}")

    # ------------------------ 网络下载 ------------------------

    def download_file(self, file_url: str, target_folder: str, file_type: Optional[str] = None) -> str:
        os.makedirs(target_folder, exist_ok=True)
        filename = os.path.basename(file_url)

        if file_type == "fwhm_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.fwhm{ext}"
        elif file_type == "preview_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.preview{ext}"

        full_path = os.path.join(target_folder, filename)
        # 已存在就跳过
        if os.path.exists(full_path) and os.path.getsize(full_path) > 0:
            return full_path

        with self.session.get(file_url, stream=True, timeout=30) as resp:
            if resp.status_code == 200:
                with open(full_path, 'wb') as f:
                    for chunk in resp.iter_content(chunk_size=1024 * 256):
                        if chunk:
                            f.write(chunk)
            else:
                err = f"Download failed: {file_url} [HTTP {resp.status_code}]"
                logger.error(err)
                resp.raise_for_status()

        return full_path

    # ------------------------ 查询参数构造 ------------------------

    def _pixscale_for_band(self, band: str) -> float:
        """
        对于 i/z 波段：
          - 2023-03-08 之前使用 Andor CCD（0.429"/pix）
          - 2023-03-08 ~ 2023-12-20 红通道没有相机（此阶段 i/z 结果可能为空）
          - 2023-12-21 起 CCD290（0.286"/pix）
        其他波段默认 0.429（可按需要调整）
        """
        pixscale = 0.429  # 默认
        if band in ["i", "z"]:
            try:
                dt_start = datetime.strptime(self.startdate, "%Y%m%d")
                if dt_start >= datetime(2023, 12, 21):
                    pixscale = 0.286
                    logger.info(f"[{band}] {dt_start.date()} 使用 CCD290，pixscale=0.286")
                elif dt_start > datetime(2023, 3, 8):
                    logger.warning(f"[{band}] {dt_start.date()} 落在无红通道相机时间段（2023-03-08 至 2023-12-20）")
                else:
                    pixscale = 0.429
                    logger.info(f"[{band}] {dt_start.date()} 使用 Andor CCD，pixscale=0.429")
            except Exception as e:
                logger.warning(f"日期解析失败，默认 pixscale=0.429：{e}")
        return pixscale

    def _build_query_params_object(self, target: str, band: str) -> Dict[str, Any]:
        encoded_target = quote(target)
        params: Dict[str, Any] = {
            "object": encoded_target,
            "filter": band,
            "start_date": self.startdate,
            "end_date": self.enddate,
            "choose": "choose",
            "au_key": self.au_key,
        }

        # 质量筛选
        if self.ell_upper_limit is not None:
            params["medEll"] = f"{self.ell_upper_limit:.2f}"
        if self.fwhm_upper_limit is not None:
            pixscale = self._pixscale_for_band(band)
            fwhm_pix = self.fwhm_upper_limit / pixscale
            params["medFWHM"] = f"{fwhm_pix:.2f}"

        return params

    def _build_query_params_radec(self, band: str) -> Dict[str, Any]:
        params: Dict[str, Any] = {
            "ra_min": self.ra_min,
            "ra_max": self.ra_max,
            "dec_min": self.dec_min,
            "dec_max": self.dec_max,
            "filter": band,
            "start_date": self.startdate,
            "end_date": self.enddate,
            "choose": "choose",
            "au_key": self.au_key,
        }

        if self.ell_upper_limit is not None:
            params["medEll"] = f"{self.ell_upper_limit:.2f}"
        if self.fwhm_upper_limit is not None:
            pixscale = self._pixscale_for_band(band)
            fwhm_pix = self.fwhm_upper_limit / pixscale
            params["medFWHM"] = f"{fwhm_pix:.2f}"

        return params

    # ------------------------ 查询+下载 主流程 ------------------------

    def _handle_item(self, item: Dict[str, Any], localpath: str, log_file_path: str):
        """处理单条记录（下载/记录/本地映射）"""
        if item.get("status") != "1":
            logger.info(f"this item {item.get('sciimg', '<unknown>')} is failed")
            return

        # focus 过滤
        sciimg_url = item.get("sciimg", "")
        calmag_url = item.get("sciimg_gaia_sexfits", "").replace("_gaia", "_CalMag")

        if not self.focus:
            if ("focus" in sciimg_url.lower()) or ("focus" in calmag_url.lower()):
                return

        if self.local:
            transformed = self.transform_urls(item)
            sciimg_local = transformed.get("sciimg", "")
            if sciimg_local:
                # 可选：检查文件是否存在
                os.system(f"stat '{sciimg_local}' >/dev/null 2>&1 || true")
                with open(log_file_path, "a") as log_file:
                    log_file.write(sciimg_local + "\n")
            return

        # 远程下载
        try:
            # 先下载 CalMag，再 sciimg
            if calmag_url:
                dst = self.download_file(calmag_url, localpath)
                # 可选：尝试生成 region（与你原脚本一致）
                try:
                    from tools import utils
                    utils.CatalogueRegion(dst)
                except Exception:
                    pass

            if sciimg_url:
                self.download_file(sciimg_url, localpath)

            # log 记录基于 ntpath.basename（与你原逻辑一致）
            with open(log_file_path, "a") as log_file:
                log_file.write(ntpath.basename(sciimg_url) + "\n")
        except Exception as e:
            logger.error(f"download error: {e}")
            with open("failed_targets.txt", "a") as fail_log:
                # 记录一个可辨识的键
                fail_log.write(f"{sciimg_url or '<no_sciimg_url>'}\n")

    # ===== object 模式 =====

    def query_and_download_object_single_target(self, target: str):
        localpath = os.path.join("images", target)
        os.makedirs(localpath, exist_ok=True)

        for band in self.bands:
            params = self._build_query_params_object(target, band)
            logger.info(f"[object] Querying [{target}] in band [{band}] ...")

            try:
                resp = self.session.get(self.base_url, params=params, timeout=30)
                resp.raise_for_status()
                data = resp.json()
            except Exception as e:
                logger.error(f"Failed to query {target} [{band}]: {e}")
                continue

            if data.get("code") != "200" or "data" not in data:
                logger.warning(f"No results for {target} [{band}]")
                continue

            log_file_path = os.path.join(localpath, f"images_{band}.list")
            with open(log_file_path, "w") as log_file:
                log_file.write("")

            for item in tqdm(data["data"], desc=f"{target}-{band}"):
                self._handle_item(item, localpath, log_file_path)

    def run_object_mode(self):
        # targetname：可以是具体目标或一个文件
        tn = self.targetname
        if tn is None:
            raise ValueError("object 模式下必须提供 --targetname")

        if os.path.isfile(tn):
            with open(tn, 'r', encoding="utf-8") as f:
                targets = [line.strip() for line in f if line.strip()]
            logger.info(f"Loaded {len(targets)} targets from file.")
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
                futures = [executor.submit(self.query_and_download_object_single_target, t) for t in targets]
                for fu in concurrent.futures.as_completed(futures):
                    fu.result()
        else:
            self.query_and_download_object_single_target(tn)

    # ===== radec 模式 =====

    def query_and_download_radec(self):
        # 目录名包含边界信息，避免覆盖
        box_name = f"RA{self.ra_min}-{self.ra_max}_DEC{self.dec_min}-{self.dec_max}"
        safe_box = re.sub(r"[^0-9A-Za-z_\-\.]", "_", box_name)
        localpath = os.path.join("images", f"radec_{safe_box}")
        os.makedirs(localpath, exist_ok=True)

        for band in self.bands:
            params = self._build_query_params_radec(band)
            logger.info(f"[radec] Querying box [{box_name}] in band [{band}] ...")

            try:
                resp = self.session.get(self.base_url, params=params, timeout=30)
                resp.raise_for_status()
                data = resp.json()
            except Exception as e:
                logger.error(f"Failed to query RA/DEC [{band}]: {e}")
                continue

            if data.get("code") != "200" or "data" not in data:
                logger.warning(f"No results in RA/DEC box for [{band}]")
                continue

            log_file_path = os.path.join(localpath, f"images_{band}.list")
            with open(log_file_path, "w") as log_file:
                log_file.write("")

            for item in tqdm(data["data"], desc=f"RADEC-{band}"):
                self._handle_item(item, localpath, log_file_path)

    # ===== 顶层 run =====

    def run(self):
        t_start = time.time()
        logger.info(f"开始处理：mode={self.mode} | api={self.api_version} | local={self.local}")

        if self.mode == "object":
            self.run_object_mode()
        else:
            self.query_and_download_radec()

        elapsed = time.time() - t_start
        logger.info(f"完成，用时：{elapsed:.2f} 秒（≈ {elapsed/60:.2f} 分钟）")


# ========================= CLI 入口 =========================

def build_parser() -> ap.ArgumentParser:
    parser = ap.ArgumentParser(description="Image Downloader for Mephisto Products (object or RA/DEC box)")

    # 模式选择：自动判断；也可显式传 --mode（可选）
    parser.add_argument("--mode", choices=["auto", "object", "radec"], default="auto",
                        help="查询模式：auto（默认，自动判断）、object、radec")

    # 通用参数
    parser.add_argument("-sd", "--startdate", type=str, default="20250101", help="起始日期 yyyymmdd")
    parser.add_argument("-ed", "--enddate", type=str, default="20250801", help="终止日期 yyyymmdd")
    parser.add_argument("-f", "--focus", action="store_true", help="是否包含 focus 图像（默认不包含）")
    parser.add_argument("-el", "--ell_upper_limit", type=float, default=None, help="椭率上限")
    parser.add_argument("-fw", "--fwhm_upper_limit", type=float, default=None, help="FWHM 上限（单位：arcsec）")
    parser.add_argument("--workers", type=int, default=4, help="最大线程数（仅 object+目标列表时有并发意义）")
    parser.add_argument("--api_version", type=str, default="V20250303_KM", help="接口版本（例如 V20250303_KM）")
    parser.add_argument("--local", action="store_true", help="local 模式：不下载，通过URL映射写出本地磁盘路径")
    parser.add_argument("--au_key", type=str, default="", help="授权Key/Token（如果接口校验需要）")
    parser.add_argument("--bands", type=str, default="u,v,g,r,i,z", help="需要的波段，逗号分隔（默认 u,v,g,r,i,z）")

    # object 模式
    parser.add_argument("-tn", "--targetname", type=str, default=None, help="目标名或目标列表文件路径")

    # radec 模式
    parser.add_argument("--ra_min", type=float, help="RA最小值 (deg)")
    parser.add_argument("--ra_max", type=float, help="RA最大值 (deg)")
    parser.add_argument("--dec_min", type=float, help="DEC最小值 (deg)")
    parser.add_argument("--dec_max", type=float, help="DEC最大值 (deg)")

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    # 模式解析
    mode = args.mode
    has_radec = all(v is not None for v in [args.ra_min, args.ra_max, args.dec_min, args.dec_max])
    if mode == "auto":
        if has_radec:
            mode = "radec"
        elif args.targetname:
            mode = "object"
        else:
            parser.error("auto 模式需要：要么提供 --targetname，要么提供 --ra_min/--ra_max/--dec_min/--dec_max")

    bands = [b.strip() for b in args.bands.split(",") if b.strip()]

    downloader = ImageDownloader(
        # 通用
        startdate=args.startdate,
        enddate=args.enddate,
        focus=args.focus,
        ell_upper_limit=args.ell_upper_limit,
        fwhm_upper_limit=args.fwhm_upper_limit,
        workers=args.workers,
        api_version=args.api_version,
        local=args.local,
        au_key=args.au_key,
        bands=bands,
        # object
        targetname=args.targetname if mode == "object" else None,
        # radec
        ra_min=args.ra_min if mode == "radec" else None,
        ra_max=args.ra_max if mode == "radec" else None,
        dec_min=args.dec_min if mode == "radec" else None,
        dec_max=args.dec_max if mode == "radec" else None,
    )

    downloader.run()


if __name__ == '__main__':
    main()


# ========================= 使用说明 =========================
"""
用法示例
--------

1) 按目标名查询（单目标）：
   python main_download_demo_img.py \
       -tn grb250419a \
       -sd 20250101 -ed 20250801 \
       --api_version V20250303_KM \
       -fw 2.5 -el 0.3

2) 按目标名查询（目标列表，每行一个目标），并发 6：
   python main_download_demo_img.py \
       -tn targets.txt \
       --workers 6 \
       -sd 20250101 -ed 20250801 \
       -fw 2.5 -el 0.3

3) RA/DEC 框选查询：
   python main_download_demo_img.py \
       --mode radec \
       --ra_min 120 --ra_max 140 \
       --dec_min 30  --dec_max 50  \
       -sd 20250101 -ed 20250801 \
       -fw 2.0 -el 0.3

4) 自动模式（推荐）：同时支持 object/RADEC，脚本自动判断：
   # object
   python main_download_demo_img.py -tn egn -sd 20250101 -ed 20250801
   # radec
   python main_download_demo_img.py --ra_min 120 --ra_max 140 --dec_min 30 --dec_max 50

其它参数
--------
- --local：不下载、仅把远端URL映射为本地磁盘路径并写入 *.list（适合在昆明/丽江内网直接访问数据盘）
- --bands：限制波段，例如 "--bands g,r,i"
- --focus：包含 focus 图像（默认不包含）
- -fw/-el：图像质量筛选。-fw 单位为角秒，将按波段/日期转换为像素（i/z 会在 2023-12-21 起使用 0.286"/pix）
"""
